#include "mvsse.h"
#include <set>
#include <vector>
#include <algorithm>

const int RSA_LEN = 2048;

void Setup(int lambda, json& ownerJson, json& usersJson, json& serverJson) {
    using nlohmann::json;
    using std::string;
    usersJson.clear();
    serverJson.clear();
    ownerJson["g"].clear();
    ownerJson["d"].clear();
    const int BUFFER_SIZE = 1024;
    char buf[BUFFER_SIZE];
    pairing_t pairing;
    pairing_init_set_str(pairing, ownerJson["param_string"].get<string>().c_str());
    element_t g, alpha, Y, s, t, K, temp_ele, temp_ele_1, tmp_Zr_ele;
    element_init_G1(g, pairing);
    element_init_G1(s, pairing);
    element_init_G1(temp_ele, pairing);
    element_init_G1(temp_ele_1, pairing);
    element_random(g);
    element_init_Zr(alpha, pairing);
    element_init_Zr(tmp_Zr_ele, pairing);
    element_random(alpha);
    int b = ownerJson["b"].get<int>();
    element_snprint(buf, BUFFER_SIZE, g);
    ownerJson["g"].push_back(buf);    // 十进制
    element_pp_t g_pp;  // 多次指数操作前的加速预处理
    element_pp_init(g_pp, g);
    element_set(tmp_Zr_ele, alpha);
    element_init_Zr(Y, pairing);    // 初始化 Y
    element_random(Y);
    element_snprint(buf, BUFFER_SIZE, Y);
    ownerJson["Y"] = buf;
    element_pp_pow_zn(s, Y, g_pp);
    element_snprint(buf, BUFFER_SIZE, s);
    ownerJson["s"] = buf;
    for (int i = 1; i <= b * 2; i++) {  // 初始化 gi
        element_pp_pow_zn(temp_ele, tmp_Zr_ele, g_pp);
        element_snprint(buf, BUFFER_SIZE, temp_ele);
        ownerJson["g"].push_back(buf);
        if (i <= b) {
            element_pow_zn(temp_ele_1, temp_ele, Y);    // 初始化 di
            element_snprint(buf, BUFFER_SIZE, temp_ele_1);
            ownerJson["d"].push_back(buf);
        }
        element_mul(tmp_Zr_ele, tmp_Zr_ele, alpha);
    }
    element_init_Zr(t, pairing);   // 随机生成 t
    element_random(t);
    element_init_GT(K, pairing);
    element_set_str(temp_ele, ownerJson["g"][b + 1].get<string>().c_str(), 10);
    pairing_apply(K, g, temp_ele, pairing); // 计算 K
    element_pow_zn(K, K, t);
    element_snprint(buf, BUFFER_SIZE, K);
    ownerJson["K"] = buf;    // 保存 K
    element_pow_zn(temp_ele, g, t); // 计算 Hdr - c0
    element_snprint(buf, BUFFER_SIZE, temp_ele);
    ownerJson["Hdr_c0"] = buf;    // 保存 Hdr - c0
    element_set(temp_ele, s);   // 计算 Hdr - c1
    json& S = ownerJson["S"];
    for (auto it = S.begin(); it != S.end(); it++) {
        int j = (*it).get<int>();
        element_set_str(temp_ele_1, ownerJson["g"][b + 1 - j].get<string>().c_str(), 10);
        element_mul(temp_ele, temp_ele, temp_ele_1);
    }
    element_pow_zn(temp_ele, temp_ele, t);
    element_snprint(buf, BUFFER_SIZE, temp_ele);
    ownerJson["Hdr_c1"] = buf;   // 保存 Hdr - c1

    // 设置 累加器 和 PRF 密钥
    gmp_randstate_t grt;
    gmp_randinit_default(grt); //设置随机数生成算法为默认
    gmp_randseed_ui(grt, time(NULL));	//设置随机化种子为当前时间
    mpz_t key;
    mpz_init(key);
    mpz_urandomb(key, grt, 128);
    ownerJson["PRF_k0"] = mpz_class(key).get_str();
    mpz_urandomb(key, grt, 128);
    ownerJson["PRF_k1"] = mpz_class(key).get_str();
    mpz_urandomb(key, grt, 128);
    ownerJson["PRF_k2"] = mpz_class(key).get_str();
    mpz_urandomb(key, grt, 128);
    ownerJson["PRF_k3"] = mpz_class(key).get_str();
    mpz_urandomb(key, grt, 128);
    ownerJson["SKE_k"] = mpz_class(key).get_str();

    // 设置累加器参数
    mpz_t P_, Q_, P, Q, N, v; // P' Q'
    mpz_init(P_);
    mpz_init(Q_);
    mpz_init(P);
    mpz_init(Q);
    mpz_init(N);
    mpz_init(v);
    mpz_urandomb(P_, grt, RSA_LEN / 2);		
    mpz_urandomb(Q_, grt, RSA_LEN / 2);
    mpz_set_si(P, 1);   
    mpz_set_si(Q, 1);
    mpz_addmul_ui(P, P_, 2);    // P += 2 * P_
    mpz_addmul_ui(Q, Q_, 2);    // Q += 2 * Q_
    mpz_mul(N, P, Q);   // N = P * Q
    // 求生成元？？
    mpz_urandomb(key, grt, RSA_LEN / 2);
    mpz_nextprime(v, key); 
    ownerJson["P"] = mpz_class(P).get_str();
    ownerJson["Q"] = mpz_class(Q).get_str();
    ownerJson["N"] = mpz_class(N).get_str();
    ownerJson["v"] = mpz_class(v).get_str();

    // 回收资源
    element_pp_clear(g_pp);
    element_clear(g);
    element_clear(alpha);
    element_clear(Y);
    element_clear(s);
    element_clear(t);
    element_clear(K);
    element_clear(temp_ele);
    element_clear(temp_ele_1);
    element_clear(tmp_Zr_ele);
    mpz_clear(key);
    mpz_clear(P_);
    mpz_clear(Q_);
    mpz_clear(P);
    mpz_clear(Q);
    mpz_clear(N);
    mpz_clear(v);
    pairing_clear(pairing);
}

void Store(json& ownerJson, json& usersJson, json& serverJson) {
    using namespace nlohmann;
    using std::string;
    using std::vector;
    json& W = ownerJson["W"];
    W.clear();
    json& D = ownerJson["D"];
    json& index = ownerJson["index"];
    index["e"].clear();

    // 构建 W
    std::set<string> wset;
    for (auto d_it = D.begin(); d_it != D.end(); ++d_it) {
        json& keywords = (*d_it)["keywords"];
        for (auto keyword_it = keywords.begin(); keyword_it != keywords.end(); ++keyword_it) {
            wset.insert((*keyword_it).get<string>());
        }
    }
    for (auto it = wset.begin(); it != wset.end(); ++it) {
        W.push_back(*it);
    }
    // 构建 index
    auto m = wset.size();
    auto n = D.size();
    index["i"] = m;
    index["j"] = m;
    index["k"] = n;
    vector<int> e(m * m * n, 0);
    vector<string> wset_vec(m);
    std::copy(wset.begin(), wset.end(), wset_vec.begin());
    for (long unsigned int k = 0; k < n; ++k) {
        json& keywords = D[k]["keywords"];
        for (long unsigned int i = 0; i < m; ++i) {
            for (long unsigned int j = i + 1; j < m; ++j) {
// std::cout << k << wset_vec[i] << wset_vec[j] << ArrayContains(keywords, wset_vec[i])<< std::endl;
                if (ArrayContains(keywords, wset_vec[i]) && ArrayContains(keywords, wset_vec[j])) {
                    e[i * m * n + j * n + k] = 1;
                }
            }
        }
    }
    for (auto it = e.begin(); it != e.end(); ++it) {
        index["e"].push_back(*it);
    }
    // 计算文件加密
    mpz_class ke(ownerJson["SKE_k"].get<string>());
    for (auto it = D.begin(); it != D.end(); ++it) {
        (*it)["C"] = SKE_Enc(ke, (*it)["text"].get<string>());
    }
    vector<mpz_class> token;
    vector<mpz_class> r;
    vector<string> label;
    
    auto& prf1 = Singleton<PRF_1>::Instance();  // 设置 PRF
    auto& prf2 = Singleton<PRF_2>::Instance();
    auto& prf3 = Singleton<PRF_3>::Instance();
    pairing_t pairing;
    pairing_init_set_str(pairing, ownerJson["param_string"].get<string>().c_str());
    element_t K, tmp_GT;
    element_init_GT(K, pairing);
    element_init_GT(tmp_GT, pairing);
    element_set_str(K, ownerJson["K"].get<string>().c_str(), 10);   // 获取 K

    char buf[BUFFER_SIZE];
    for (long unsigned int i = 0; i < m; ++i) {
        for (long unsigned int j = i + 1; j < m; ++j) {
            string hex_str = prf1(wset_vec[i] + wset_vec[j]);
            token.emplace_back(hex_str, 16);
            r.emplace_back(prf3(hex_str), 16);
            element_pow_mpz(tmp_GT, K, (token.back() * r.back()).get_mpz_t());
            element_snprint(buf, BUFFER_SIZE, tmp_GT);
            label.push_back(buf);
        }
    }


}


